"""前缀树(字典树)"""


class Trie:
    class Node:
        def __init__(self, is_word=False):
            self.is_word = is_word
            self.next_node = {}

    def __init__(self):
        self.root = Trie.Node()
        self._size = 0

    def size(self):
        return self._size

    def __len__(self):
        return self._size

    def add(self, word):
        """向前缀树中添加一个单次"""
        current = self.root
        for c in word:
            # if c not in current.next_node:
            #     current.next_node[c] = Trie.Node()
            current.next_node.setdefault(c, Trie.Node())
            current = current.next_node[c]
        # 到达单次的末尾
        if not current.is_word:  # 如果单次已经存在了，就不把size增加1
            current.is_word = True
            self._size += 1

    def __contains__(self, word):
        current = self.root
        for c in word:
            if c not in current.next_node:
                return False
            current = current.next_node[c]
        if current.is_word:  # 包含了所有的字母不代表包含这个单次，包不包含还要看is_word
            return True
        else:
            return False

    def contains(self, word):
        """查看前缀树中是否包含某个单次"""
        current = self.root
        for c in word:
            if c not in current.next_node:
                return False
            current = current.next_node[c]
        return current.is_word

    def is_prefix(self, prefix):
        """判断前缀树中是否有某个前缀"""
        current = self.root
        for c in prefix:
            if c not in current.next_node:
                return False
            current = current.next_node[c]
        return True
